﻿using System;
using System.Linq;
using System.Collections.Generic;
using System.ComponentModel;
using System.Diagnostics;
using System.Runtime.InteropServices;

namespace AshuSoft
{
    // ReSharper disable EnumUnderlyingTypeIsInt
    public enum ThreadAccess : int

    {
      TERMINATE = (0x0001),
      SUSPEND_RESUME = (0x0002),
      GET_CONTEXT = (0x0008),
      SET_CONTEXT = (0x0010),
      SET_INFORMATION = (0x0020),
      QUERY_INFORMATION = (0x0040),
      SET_THREAD_TOKEN = (0x0080),
      IMPERSONATE = (0x0100),
      DIRECT_IMPERSONATION = (0x0200)
    }
    // ReSharper restore EnumUnderlyingTypeIsInt
    
    /// <summary>
    /// A utility class to determine a process parent.
    /// </summary>
    [StructLayout(LayoutKind.Sequential)]
    public struct ParentProcessUtilities
    {
        // These members must match PROCESS_BASIC_INFORMATION
        internal IntPtr Reserved1;
        internal IntPtr PebBaseAddress;
        internal IntPtr Reserved2_0;
        internal IntPtr Reserved2_1;
        internal IntPtr UniqueProcessId;
        internal IntPtr InheritedFromUniqueProcessId;

        [DllImport("ntdll.dll")]
        private static extern int NtQueryInformationProcess(IntPtr processHandle, int processInformationClass, ref ParentProcessUtilities processInformation, int processInformationLength, out int returnLength);

        #region Suspension methods
        [DllImport("kernel32.dll")]
        static extern IntPtr OpenThread(ThreadAccess dwDesiredAccess, bool bInheritHandle, uint dwThreadId);
        [DllImport("kernel32.dll")]
        static extern uint SuspendThread(IntPtr hThread);
        [DllImport("kernel32.dll")]
        static extern int ResumeThread(IntPtr hThread);

        private static void SuspendProcess(int processId)
        {
          Process proc = Process.GetProcessById(processId);

          if (proc.ProcessName == string.Empty)
            return;

            foreach (ProcessThread pT in proc.Threads)
            {
                IntPtr pOpenThread = OpenThread(ThreadAccess.SUSPEND_RESUME, false, (uint)pT.Id);

                if (pOpenThread == IntPtr.Zero)
                {
                break;
                }

                SuspendThread(pOpenThread);
            }
        }

        public static void ResumeProcess(int processId)
        {
            Process proc = Process.GetProcessById(processId);

            if (proc.ProcessName == string.Empty)
                return;

            foreach (ProcessThread pT in proc.Threads)
            {
                IntPtr pOpenThread = OpenThread(ThreadAccess.SUSPEND_RESUME, false, (uint)pT.Id);

                if (pOpenThread == IntPtr.Zero)
                {
                    break;
                }

                ResumeThread(pOpenThread);
            }
        }
        #endregion
        
        private static bool IsParentProcess(int parentProcessId, int id)
        {
            try
            {
                Process process = Process.GetProcessById(id);
                return IsParentProcess(parentProcessId, process.Handle);
            }
            catch(ArgumentException){}
            catch(InvalidOperationException){}
            catch (Win32Exception){}
            
            return false;
        }

        private static bool IsParentProcess(int parentProcessId, IntPtr handle)
        {
            var pbi = new ParentProcessUtilities();
            int returnLength;
            int status = NtQueryInformationProcess(handle, 0, ref pbi, Marshal.SizeOf(pbi), out returnLength);
            if (status != 0)
                return false;
            //throw new Win32Exception(status);

            try
            {
                return parentProcessId == pbi.InheritedFromUniqueProcessId.ToInt32();
            }
            catch (ArgumentException)
            {
                // not found
                return false;
            }
        }

        public static List<Process> GetChildProcess(int parentId, string processNameFilter)
        {
            var childProcesses = new List<Process>();
            foreach (Process process in Process.GetProcesses())
            {
                string processName = process.ProcessName ?? "";
                bool nameMatch = true;
                if (!String.IsNullOrWhiteSpace(processName) && !String.IsNullOrWhiteSpace(processNameFilter))
                {
                    nameMatch = processName.StartsWith(processNameFilter, StringComparison.OrdinalIgnoreCase);
                }
                if(process.Id != parentId && IsParentProcess(parentId, process.Id) && nameMatch)
                {
                    childProcesses.Add(process);
                }
            }
            return childProcesses;
        }

        public static List<Process> GetChildProcessTree(int parentId, string immediateChildNameFilter)
        {
            var childProcesses = new List<Process>();
            childProcesses.AddRange(GetChildProcess(parentId, immediateChildNameFilter));
            int i = 0;
            while(i < childProcesses.Count)
            {
                var process = childProcesses[i];
                childProcesses.AddRange(GetChildProcess(process.Id, null));
                i++;
            }
            return childProcesses;
        }

        public static void KillProcessChild(int processId, string immediateChildNameFilter)
        {
            List<Process> processToKill = new List<Process>();

            GetAndSuspendChild(processId, processToKill, immediateChildNameFilter);

            foreach (var process in processToKill)
            {
                Exception ex = null;
                try
                {
                    process.Kill();
                    Trace.WriteLine(String.Format("Killed Process {0}:{1}", process.ProcessName, process.Id));
                }
                catch(Win32Exception win32Exception)
                {
                    ex = win32Exception;
                }
                catch(NotSupportedException notSupportedException)
                {
                    ex = notSupportedException;
                }
                catch(InvalidOperationException invalidOperationException)
                {
                    ex = invalidOperationException;
                }
                if(ex != null)
                {
                    Trace.WriteLine(String.Format("Failed to kill Process {0}:{1}", process.ProcessName, process.Id));
                }
            }
        }

        private static int _currentProcessId = -1;
        static int CurrentProcessId 
        {
            get 
            {
                if (_currentProcessId != -1) return _currentProcessId;
                _currentProcessId = Process.GetCurrentProcess().Id;
                return _currentProcessId;
            }
        }

        static void GetAndSuspendChild(int processId, List<Process> processToKill, string immediateChildNameFilter)
        {
            var child = GetChildProcessTree(processId, immediateChildNameFilter);
            if(child == null || child.Count == 0)return;
            var currentProcess = child.FirstOrDefault(p => p.Id == CurrentProcessId);
            if (currentProcess != null && child.Contains(currentProcess)) 
                child.Remove(currentProcess);

            processToKill.AddRange(child);

            foreach (var process in child.ToArray()) 
            {
                try
                {
                    SuspendProcess(process.Id);
                    Trace.WriteLine(String.Format("Suspended Process {0}:{1}", process.ProcessName, process.Id));
                }
                catch(ArgumentException)
                {
                    Trace.WriteLine(String.Format("Process {0}:{1} probably already exited.", process.ProcessName, process.Id));
                    //child.Remove(process);
                }
            }
            foreach (var process in child)
            {
                GetAndSuspendChild(process.Id, processToKill, null);
            }
        }
    }

}
